import re
import paddle
from paddlenlp.transformers import BertTokenizer
from PoetryGen import PoetryGen
from Poetry_Bert_Model import PoetryBertModel
bert_tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
net = PoetryBertModel('bert-base-chinese', 128)
model = paddle.Model(net)
model.load('./checkpoint_10/final.pdopt')
poetry_gen = PoetryGen(model, bert_tokenizer)
# 遇到标点符号换行的工具函数
def poetry_show(poetry):
    pattern = r"([，。；？])"
    text = re.sub(pattern, r'\1 ', poetry)
    for p in text.split():
        if p:
            print(p)

while True:
    text = input("开始吧：")
    # 续写
    # style其实也作为历史放入其中，作为提示
    poetry = poetry_gen.generate(head=text,mode=2,style="")
    poetry_show(poetry)



